#lang scheme

; TODO: ipmask
; ip4 = (bytevector 4)   ; TODO: define-blob 'IP4 4
; ip16 = (bytevector 16) ; TODO: define-blob 'IP16 16
(provide string->ip  string->ip4  string->ip6  ; convert string to packed ip
         bytes->ip   bytes->ip4   bytes->ip6   ; convert string to packed ip
         string-ip?  string-ip4?  string-ip6?  ; is string represent ip?
         bytes-ip?   bytes-ip4?   bytes-ip6?   ; is bytestring represent ip?
         ip4? ip6?
         ip->string  ip4->string  ip6->string
         ip->bytes   ip4->bytes   ip6->bytes
         ip->uint    ip4->uint32  ip6->uint128
         ip-compare  ip4-compare  ip6-compare
         )

(require "test.ss")
(require "bytes.ss")
(require scheme/mpair) ; mz bug
(require srfi/1) ; for every
(require srfi/13) ; for string-index
(require srfi/26) ; for cut
(require (except-in srfi/67 string-compare-ci string-compare)) ; if3
(require rnrs/bytevectors-6)


(define (ip4? ip)
  (= 4 (bytevector-length ip)))

(define (ip6? ip)
  (= 16 (bytevector-length ip)))


(define (string-ip? str)
  (or (string-ip4? str)
      (string-ip6? str)))

(define (bytes-ip? str)
  (or (bytes-ip4? str)
      (bytes-ip6? str)))

(define (string-ip4? str)
  (let ([parts (string-split str #\.)]) ; todo: . at the end should be ok
    (and (= 4 (length parts))
         (every (lambda (part) 
                  (and (every char-numeric? (string->list part))
                       (<= 0 (string->number part) 255))) parts))))

(define (string-ip6? str)
  (error "string-ip6? not implemented yet"))

(define (bytes-ip4? str)
  (let ([parts (bytes-split str (char->integer #\.))]) ; todo: . at the end should be ok
    (and (= 4 (length parts))
         (every (lambda (part) 
                  (and (every byte-numeric? (bytes->list part))
                       (<= 0 (bytes->number part) 255))) parts))))

(define (bytes-ip6? str)
  (error "bytes-ip6? not implemented yet"))


(define (string->ip str)
  ((if (string-index str #\.) string->ip4 string->ip6) str))

(define (string->ip4 str)
  (let ([parts (string-split str #\.)]) ; todo: . at the end should be ok
    (unless (= 4 (length parts))
      (error "string->ip: invalid ip" str))
    (list->bytes (map string->number parts)))) ; mz bug?: u8-list->bytevector expects mutable list 

(define (string->ip6 str)
  (let ([parts (string-split str #\:)]) ; todo: support short forms
    (unless (= 8 (length parts))
      (error "string->ip: invalid ip" str))
    (uint-list->bytevector
     (list->mlist ; mz bug
      (map (cut string->number <> 16) parts))
     (endianness big) 2)))


(define (bytes->ip bstr)
  ((if (bytes-index bstr (char->integer #\.)) bytes->ip4 bytes->ip6) bstr))

(define (bytes->ip4 bstr)
  (let ([parts (bytes-split bstr (char->integer #\.))]) ; todo: . at the end should be ok
    (unless (= 4 (length parts))
      (error "bytes->ip: invalid ip" bstr))
    (list->bytes (map bytes->number parts)))) ; mz bug?: u8-list->bytevector expects mutable list 

(define (bytes->ip6 bstr)
  (let* ([parts2 (bytes-split bstr (char->integer #\:))]
         [parts (bytes-split bstr (char->integer #\:))]
         [missed (- 8 (length parts))])
    (when (< missed 0)
      (error "bytes->ip6: invalid ip" bstr))
    
    ; replace :: with missed zeroes
    ; TODO: make it on vector
    (let ([parts-fixed 
           (foldr (lambda (e l) 
                    (if (bytes-null? e)
                        (if (or (null? l) ; first and last
                                (= 7 (length l)))
                            (cons 0 l)
                            (case missed
                              [(0) (cons  0 l)]
                              [(1) (cons* 0 0 l)]
                              [(2) (cons* 0 0 0 l)]
                              [(3) (cons* 0 0 0 0 l)]
                              [(4) (cons* 0 0 0 0 0 l)]
                              [(5) (cons* 0 0 0 0 0 0 l)]
                              [(6) (cons* 0 0 0 0 0 0 0 l)]
                              [(7) (cons* 0 0 0 0 0 0 0 0 l)]
                              [(8) (cons* 0 0 0 0 0 0 0 0 0 l)]))
                        (cons (bytes->number e 16) l))) 
                  '() parts)])
;      (o 'FIXED parts-fixed)
      
      (uint-list->bytevector
       (list->mlist ; mz bug
        parts-fixed)
       (endianness big) 2))))


(define (ip->string ip)
  (cond 
    [(ip4? ip) (ip4->string ip)]
    [(ip6? ip) (ip6->string ip)]))

(define (ip4->string ip4)
  (string-append
   (number->string (bytevector-u8-ref ip4 0)) "."
   (number->string (bytevector-u8-ref ip4 1)) "."
   (number->string (bytevector-u8-ref ip4 2)) "."
   (number->string (bytevector-u8-ref ip4 3))))

; TODO: optimise for shorted addeses
(define (ip6->string ip6 [full #f])
  (string-append
   (number->string (bytevector-u16-ref ip6 0 (endianness big)) 16) ":"
   (number->string (bytevector-u16-ref ip6 2 (endianness big)) 16) ":"
   (number->string (bytevector-u16-ref ip6 4 (endianness big)) 16) ":"
   (number->string (bytevector-u16-ref ip6 6 (endianness big)) 16) ":"
   (number->string (bytevector-u16-ref ip6 8 (endianness big)) 16) ":"
   (number->string (bytevector-u16-ref ip6 10 (endianness big)) 16) ":"
   (number->string (bytevector-u16-ref ip6 12 (endianness big)) 16) ":"
   (number->string (bytevector-u16-ref ip6 14 (endianness big)) 16)))


(define (ip->bytes ip)
  (cond 
    [(ip4? ip) (ip4->bytes ip)]
    [(ip6? ip) (ip6->bytes ip)]))

(define (ip4->bytes ip4)
  (bytes-append
   (number->bytes (bytevector-u8-ref ip4 0)) #"."
   (number->bytes (bytevector-u8-ref ip4 1)) #"."
   (number->bytes (bytevector-u8-ref ip4 2)) #"."
   (number->bytes (bytevector-u8-ref ip4 3))))

; TODO: optimise for shorted addeses
(define (ip6->bytes ip6 [full #f])
  (bytes-append
   (number->bytes (bytevector-u16-ref ip6 0 (endianness big)) 16) #":"
   (number->bytes (bytevector-u16-ref ip6 2 (endianness big)) 16) #":"
   (number->bytes (bytevector-u16-ref ip6 4 (endianness big)) 16) #":"
   (number->bytes (bytevector-u16-ref ip6 6 (endianness big)) 16) #":"
   (number->bytes (bytevector-u16-ref ip6 8 (endianness big)) 16) #":"
   (number->bytes (bytevector-u16-ref ip6 10 (endianness big)) 16) #":"
   (number->bytes (bytevector-u16-ref ip6 12 (endianness big)) 16) #":"
   (number->bytes (bytevector-u16-ref ip6 14 (endianness big)) 16)))



(define (ip->uint ip)
  (cond 
    [(ip4? ip) (ip4->uint32 ip)]
    [(ip6? ip) (ip6->uint128 ip)]))

; TODO: bytevector-u32
(define (ip4->uint32 ip4)
  (bitwise-ior 
   (arithmetic-shift (vector-ref ip4 0) 24)
   (arithmetic-shift (vector-ref ip4 1) 16)
   (arithmetic-shift (vector-ref ip4 2) 8)
   (vector-ref ip4 3)))

; TODO: bytevector-uint
(define (ip6->uint128 ip6)
  (bitwise-ior 
   (arithmetic-shift (vector-ref ip6 0) 120)
   (arithmetic-shift (vector-ref ip6 1) 112)
   (arithmetic-shift (vector-ref ip6 2) 104)
   (arithmetic-shift (vector-ref ip6 3) 96)
   (arithmetic-shift (vector-ref ip6 4) 88)
   (arithmetic-shift (vector-ref ip6 5) 80)
   (arithmetic-shift (vector-ref ip6 6) 72)
   (arithmetic-shift (vector-ref ip6 7) 64)
   (arithmetic-shift (vector-ref ip6 8) 56)
   (arithmetic-shift (vector-ref ip6 9) 48)
   (arithmetic-shift (vector-ref ip6 10) 40)
   (arithmetic-shift (vector-ref ip6 11) 32)
   (arithmetic-shift (vector-ref ip6 12) 24)
   (arithmetic-shift (vector-ref ip6 13) 16)
   (arithmetic-shift (vector-ref ip6 14) 8)
   (vector-ref ip6 15)))

; assume any ip4 < any ip6
(define (ip-compare a b)
  (if3 (integer-compare (bytes-length a) (bytes-length b))
       -1
       (bytes-compare a b)
       1))

(define ip4-compare bytes-compare)

(define ip6-compare bytes-compare)

; TESTS
(=! -1 (ip-compare (string->ip "192.168.1.20") (string->ip6 "1111:2222:3333:4444:5555:6666:7777:8888")))
(=! 0 (ip-compare (string->ip "192.168.1.20") (string->ip "192.168.1.20")))
(=! 1 (ip-compare (string->ip "192.168.1.21") (string->ip "192.168.1.20")))

(equal! (ip6->bytes (bytes->ip6 #"1111:2222:3333:4444:5555:6666:7777:8888") 'full) #"1111:2222:3333:4444:5555:6666:7777:8888")
(equal! (ip6->bytes (bytes->ip6 #"1111:2222:3333:4444:5555:6666::8888") 'full)     #"1111:2222:3333:4444:5555:6666:0:8888")
(equal! (ip6->bytes (bytes->ip6 #"2001:6f8:129d:1000::10") 'full)                  #"2001:6f8:129d:1000:0:0:0:10")
(equal! (ip6->bytes (bytes->ip6 #"::") 'full)                                      #"0:0:0:0:0:0:0:0")
(equal! (ip6->bytes (bytes->ip6 #"::0") 'full)                                     #"0:0:0:0:0:0:0:0")
(equal! (ip6->bytes (bytes->ip6 #"0::") 'full)                                     #"0:0:0:0:0:0:0:0")
(equal! (ip6->bytes (bytes->ip6 #"0::0") 'full)                                    #"0:0:0:0:0:0:0:0")
(equal! (ip6->bytes (bytes->ip6 #"0:1::0") 'full)                                  #"0:1:0:0:0:0:0:0")
(equal! (ip6->bytes (bytes->ip6 #"0::1:0") 'full)                                  #"0:0:0:0:0:0:1:0")
(equal! (ip6->bytes (bytes->ip6 #":1::2") 'full)                                   #"0:1:0:0:0:0:0:2")
(equal! (ip6->bytes (bytes->ip6 #"::1:") 'full)                                    #"0:0:0:0:0:0:1:0")

;(equal! (ip6->bytes (bytes->ip6 #"1111:2222:3333:4444:5555:6666:7777:8888") 'full) #"1111:2222:3333:4444:5555:6666:7777:8888")
;(equal! (ip6->bytes (bytes->ip6 #"1111:2222:3333:4444:5555:6666::8888") 'full)     #"1111:2222:3333:4444:5555:6666::8888")
;(equal! (ip6->bytes (bytes->ip6 #"2001:6f8:129d:1000::10") 'full)                  #"2001:6f8:129d:1000::10")
;(equal! (ip6->bytes (bytes->ip6 #"::") 'full)                                      #"0::0")
;(equal! (ip6->bytes (bytes->ip6 #"::0") 'full)                                     #"0::0")
;(equal! (ip6->bytes (bytes->ip6 #"0::") 'full)                                     #"0::0")
;(equal! (ip6->bytes (bytes->ip6 #"0::0") 'full)                                    #"0::0")
;(equal! (ip6->bytes (bytes->ip6 #"0:1::0") 'full)                                  #"0:1::0")
;(equal! (ip6->bytes (bytes->ip6 #"0::1:0") 'full)                                  #"0::1:0")
;(equal! (ip6->bytes (bytes->ip6 #":1::2") 'full)                                   #"0:1::2")
;(equal! (ip6->bytes (bytes->ip6 #"::1:") 'full)                                    #"0::1:0")

(equal! (string->ip "192.168.1.20") #"\300\250\1\24")
(equal! (bytes->ip #"192.168.1.20") #"\300\250\1\24")
(equal! (ip->bytes #"\300\250\1\24") #"192.168.1.20")
(equal! (ip->string #"\300\250\1\24") "192.168.1.20")

(equal! (string->ip6 "1111:2222:3333:4444:5555:6666:7777:8888") #"\21\21\"\"33DDUUffww\210\210")
(equal! (bytes->ip6 #"1111:2222:3333:4444:5555:6666:7777:8888") #"\21\21\"\"33DDUUffww\210\210")
(equal! (ip6->string #"\21\21\"\"33DDUUffww\210\210") "1111:2222:3333:4444:5555:6666:7777:8888")
(equal! (ip6->bytes #"\21\21\"\"33DDUUffww\210\210") #"1111:2222:3333:4444:5555:6666:7777:8888")

(true! (string-ip4? "192.168.1.20"))
(false! (string-ip4? "192.168.1.256")
        (string-ip4? "192.168.1.1.1")
        (string-ip4? "192.168.1"))

(true! (bytes-ip4? #"192.168.1.20"))
(false! (bytes-ip4? #"192.168.1.256")
        (bytes-ip4? #"192.168.1.1.1")
        (bytes-ip4? #"192.168.1"))